import os
import glob
import xml.etree.ElementTree as ET
import matplotlib.pyplot as plt
import yaml
import numpy as np
from PIL import Image
from collections import Counter, defaultdict
import matplotlib
matplotlib.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
matplotlib.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

def check_dataset():
    # 设置工作目录
    work_dir = os.getcwd()
    
    # 加载类别配置
    with open('labels_config.yaml', 'r') as f:
        config = yaml.safe_load(f)
    
    # 获取所有类别名称
    class_names = [v for _, v in config['names'].items()]
    
    # 获取所有XML和PNG文件
    xml_files = set(glob.glob('*.xml'))
    png_files = set(glob.glob('*.png'))
    
    # 图片总数量
    total_images = len(png_files)
    
    # 检查图片和标注文件是否一一对应
    xml_basenames = {os.path.basename(xml_file).split('.')[0] for xml_file in xml_files}
    png_basenames = {os.path.basename(png_file).split('.')[0] for png_file in png_files}
    
    xml_only = xml_basenames - png_basenames
    png_only = png_basenames - xml_basenames
    
    # 统计每个类别的图片数量和检测框数量
    class_image_count = defaultdict(set)  # 每个类别包含的图片集合
    class_bbox_count = Counter()  # 每个类别的边界框数量
    
    # 记录损坏的图片
    corrupted_images = []
    
    # 边界框总数
    total_bboxes = 0
    
    # 解析所有XML文件
    for xml_file in xml_files:
        try:
            tree = ET.parse(xml_file)
            root = tree.getroot()
            
            # 获取图片文件名
            filename = os.path.basename(xml_file).split('.')[0]
            
            # 统计每个对象的类别
            for obj in root.findall('object'):
                # 在这个数据集中，类别可能使用不同的标签名称，尝试多种可能性
                class_name = None
                
                # 尝试从'n'标签获取类别名称
                name_tag = obj.find('n')
                if name_tag is not None and name_tag.text is not None:
                    class_name = name_tag.text
                else:
                    # 尝试从'name'标签获取类别名称
                    name_tag = obj.find('name')
                    if name_tag is not None and name_tag.text is not None:
                        class_name = name_tag.text
                
                if class_name is None:
                    print(f"警告: 无法在 {xml_file} 中找到类别名称，跳过该对象")
                    continue
                
                # 记录该类别包含了这张图片
                class_image_count[class_name].add(filename)
                
                # 增加该类别的边界框计数
                class_bbox_count[class_name] += 1
                
                # 增加总边界框计数
                total_bboxes += 1
        except Exception as e:
            print(f"解析 {xml_file} 时出错: {str(e)}")
    
    # 检查图片是否损坏
    for png_file in png_files:
        try:
            img = Image.open(png_file)
            img.verify()  # 验证图像文件
        except Exception as e:
            corrupted_images.append(png_file)
            print(f"图片 {png_file} 已损坏: {str(e)}")
    
    # 计算每个类别的图片数量
    class_image_counts = {class_name: len(images) for class_name, images in class_image_count.items()}
    
    # 按类别名称排序
    sorted_class_names = sorted(class_image_counts.keys())
    sorted_image_counts = [class_image_counts.get(name, 0) for name in sorted_class_names]
    sorted_bbox_counts = [class_bbox_count.get(name, 0) for name in sorted_class_names]
    
    # 输出统计结果
    print("\n===== 数据集统计信息 =====")
    print(f"图片总数: {total_images}")
    print(f"检测框总数: {total_bboxes}")
    
    if total_images > 0:
        print(f"平均每张图片的检测框数量: {total_bboxes / total_images:.2f}")
    
    # 输出每个类别的统计信息
    print("\n类别统计:")
    for class_name in sorted_class_names:
        print(f"  {class_name}: {class_image_counts.get(class_name, 0)} 张图片, {class_bbox_count.get(class_name, 0)} 个检测框")
    
    # 打印缺失对应关系的文件
    if xml_only:
        print("\n以下XML文件没有对应的图片文件:")
        for basename in sorted(xml_only):
            print(f"  {basename}.xml")
    
    if png_only:
        print("\n以下图片文件没有对应的XML文件:")
        for basename in sorted(png_only):
            print(f"  {basename}.png")
    
    if corrupted_images:
        print("\n以下图片文件已损坏:")
        for img in corrupted_images:
            print(f"  {img}")
    
    # 创建结果目录
    results_dir = 'dataset_check_results'
    os.makedirs(results_dir, exist_ok=True)
    
    # 可视化每个类别的图片数量
    plt.figure(figsize=(15, 8))
    bars = plt.bar(sorted_class_names, sorted_image_counts, color='skyblue')
    plt.title('每个类别的图片数量', fontsize=16)
    plt.xlabel('类别名称', fontsize=14)
    plt.ylabel('图片数量', fontsize=14)
    plt.xticks(rotation=90, fontsize=10)
    plt.tight_layout()
    
    # 在柱状图上添加数值标签
    for bar in bars:
        height = bar.get_height()
        plt.text(bar.get_x() + bar.get_width()/2., height + 0.1,
                f'{int(height)}',
                ha='center', va='bottom', fontsize=9)
    
    plt.savefig(os.path.join(results_dir, 'class_image_count.png'), dpi=300)
    plt.close()
    
    # 可视化每个类别的检测框数量
    plt.figure(figsize=(15, 8))
    bars = plt.bar(sorted_class_names, sorted_bbox_counts, color='lightgreen')
    plt.title('每个类别的检测框数量', fontsize=16)
    plt.xlabel('类别名称', fontsize=14)
    plt.ylabel('检测框数量', fontsize=14)
    plt.xticks(rotation=90, fontsize=10)
    plt.tight_layout()
    
    # 在柱状图上添加数值标签
    for bar in bars:
        height = bar.get_height()
        plt.text(bar.get_x() + bar.get_width()/2., height + 0.1,
                f'{int(height)}',
                ha='center', va='bottom', fontsize=9)
    
    plt.savefig(os.path.join(results_dir, 'class_bbox_count.png'), dpi=300)
    plt.close()
    
    print(f"\n可视化结果已保存到 {results_dir} 目录")

if __name__ == '__main__':
    check_dataset() 